package com.xiashitech.agent.instrumentation.spring.controller_earlier_v3_1.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import com.xiashitech.agent.utils.TypeUtils;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Optional;

import static com.xiashitech.agent.instrumentation.spring.controller_earlier_v3_1.XSSpringControllerSingletons.controllerInstrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSSpringControllerMethodInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("org.springframework.web.servlet.DispatcherServlet")
                .and(ElementMatchers.not(hasClassesNamed("org.springframework.web.servlet.mvc.method.annotation.ServletInvocableHandlerMethod")));
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("org.springframework.web.servlet.DispatcherServlet");
    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("doService"))
                        .and(takesArgument(0, named("javax.servlet.http.HttpServletRequest")))
                        .and(takesArgument(1, named("javax.servlet.http.HttpServletResponse")))
                        .and(isProtected()),
                this.getClass().getName() + "$DoServiceAdvice");
    }


    public static class DoServiceAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(value = 0,readOnly = false) HttpServletRequest request,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("user") UserRequest userRequest) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSpringControllerMethodInstrumentation");
            try {
                if(TypeUtils.isSkipAcceptType(request.getHeader("Accept"), request.getRequestURI())) {
                    return;
                }
                if(CacheStatus.cannotRun()) {
                    return;
                }
                Optional<UserRequest> userRequestOptional = Java8Bridge.executeRuleForUserInfoMatchRule(request);
                if(userRequestOptional.isPresent()) {
                    userRequest = userRequestOptional.get();
                } else {
                    userRequest = UserRequest.builder().build();
                }
                userRequest.setRequest(request);
                Context parentContext = currentContext();
                if (controllerInstrumenter().shouldStart(parentContext,userRequest)) {
                    parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                    context = controllerInstrumenter().start(parentContext, userRequest);
                    context = BaggageUtil.syncBaggage()
                            .put("xs.userId",userRequest.getUserid() == null ? "" : userRequest.getUserid())
                            .put("xs.loginName",userRequest.getLoginName() == null ? "" : userRequest.getLoginName())
                            .put("xs.sessionId",userRequest.getSessionId() == null ? "" : userRequest.getSessionId())
                            .put("xs.requestDate",userRequest.getRequestDate() == null ? "" : userRequest.getRequestDate().toString())
                            .put("xs.sessionUpdate",userRequest.getSessionUpdate() == null ? "false" : userRequest.getSessionUpdate().toString())
                            .put("endpoint-api",userRequest.getEndpointApi())
                            .put("api-systemId",userRequest.getApiSystemId())
                            .build().storeInContext(context);
                    scope = context.makeCurrent();
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("DoServiceAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Argument(value = 0, readOnly = false) HttpServletRequest request,
                @Advice.Argument(value = 1, readOnly = false) HttpServletResponse response,
                @Advice.Thrown Throwable throwable,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("user") UserRequest userRequest) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSpringControllerMethodInstrumentation");
            try {
                if(TypeUtils.isSkipAcceptType(request.getHeader("Accept"), request.getRequestURI())) {
                    return;
                }
                if (CacheStatus.cannotRun()) {
                    return;
                }
                if (scope == null) {
                    return;
                }
                if (userRequest.getUserid() == null || "".equals(userRequest.getUserid())) {
                    Optional<UserRequest> userRequestOptional = Java8Bridge.executeRuleForUserInfoMatchRule(request);
                    if (userRequestOptional.isPresent() && userRequestOptional.get().getUserid() != null && !"".equals(userRequestOptional.get().getUserid())) {
                        userRequest = userRequestOptional.get();
                        userRequest.setSessionUpdate(true);
                    }
                }
                scope.close();
                userRequest.setRequest(request);
                controllerInstrumenter().end(context, userRequest, response, throwable);
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("DoServiceAdvice e");
        }
    }
}
